package com.example.webfluxdemo;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.channels.Channels;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.stream.Collectors;

import com.alibaba.fastjson.JSONObject;
import org.reactivestreams.Publisher;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.server.PathContainer;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.ServerWebExchangeDecorator;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import org.springframework.web.util.pattern.PathPattern;
import org.springframework.web.util.pattern.PathPatternParser;

import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

@Component
@Slf4j
public class ApiFilter implements WebFilter {

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        long start = System.currentTimeMillis();
        PathPattern[] patterns = new PathPattern[] {
            new PathPatternParser().parse("/*/favicon.icon"),
//            new PathPatternParser().parse("/error"),
        };
        PathContainer container = exchange.getRequest().getPath().pathWithinApplication();
        for (PathPattern pattern : patterns) {
            if (pattern.matches(container)) {
                return chain.filter(exchange);
            }
        }
        String requestId = UUID.randomUUID().toString().replaceAll("-", "");
        Map<String, String> elkLog = new HashMap<>();
        elkLog.put("requestId" , requestId);
        return chain.filter(new ApiWebExchange(exchange, elkLog))
                .contextWrite(ctx -> ctx.put("isDebug", true))
                .contextWrite(ctx -> ctx.put("elk", elkLog))
                .doOnSuccess(aVoid -> log.info("Elapsed: {}ms", System.currentTimeMillis() - start))
                .doOnError(aVoid -> log.info("Error Elapsed: {}ms", System.currentTimeMillis() - start));
    }

    static class ApiWebExchange extends ServerWebExchangeDecorator {
        private final ApiRequestDecorator request;
        private final ApiResponseDecorator response;

        protected ApiWebExchange(ServerWebExchange delegate, Map<String, String> elkLog) {
            super(delegate);
            request = new ApiRequestDecorator(delegate.getRequest(), elkLog);
            response = new ApiResponseDecorator(delegate.getResponse(), elkLog);
        }

        @Override
        public ServerHttpRequest getRequest() {
            return this.request;
        }
        @Override
        public ServerHttpResponse getResponse() {
            return this.response;
        }
    }

    static class ApiRequestDecorator extends ServerHttpRequestDecorator {
        private final Flux<DataBuffer> body;
        public ApiRequestDecorator(ServerHttpRequest request, Map<String, String> elkLog) {
            super(request);
            String url = request.getURI().getPath();
            if (url.contains("/event"))
                url = Arrays.stream(url.split("/")).limit(3).collect(Collectors.joining("/"));

            elkLog.put("ActionName", url);
            elkLog.put("Path", request.getPath().toString());
            elkLog.put("ip", IpUtils.getRequestIp(request));
            elkLog.put("url", String.format("%s://%s:%s%s?%s", request.getURI().getScheme(),
                    request.getURI().getHost(),
                    request.getURI().getPort(),
                    request.getURI().getPath(),
                    request.getURI().getQuery()));
            body = super.getBody().doOnNext(buffer -> {
                ByteArrayOutputStream bodyStream = new ByteArrayOutputStream();
                try {
                    Channels.newChannel(bodyStream).write(buffer.asByteBuffer().asReadOnlyBuffer());
                    elkLog.put("requestBody", bodyStream.toString());
                } catch (IOException e) {
                    log.warn("Failed to get request body: {}", e.getMessage());
                }
            });
        }

        @Override
        public Flux<DataBuffer> getBody() {
            return body;
        }
    }

    static class ApiResponseDecorator extends ServerHttpResponseDecorator {
        private final Map<String, String> elkLog;
        public ApiResponseDecorator(ServerHttpResponse delegate, Map<String, String> elkLog) {
            super(delegate);
            this.elkLog = elkLog;
        }

        @Override
        public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
            return super.writeWith(Flux.from(body).doOnNext(buffer -> {
                ByteArrayOutputStream bodyStream = new ByteArrayOutputStream();
                try {
                    Channels.newChannel(bodyStream).write(buffer.asByteBuffer().asReadOnlyBuffer());
                    elkLog.put("ResponseBody", bodyStream.toString());
                } catch (IOException e) {
                    log.error("Failed to get response body: {}", e.getMessage());
                } finally {
                    log.info("ELK: {}", JSONObject.toJSONString(elkLog));
                }
            }));
        }
    }
}
